#include "meshFunctions.hpp"
#include "meshTypeFunctions.hpp"
#include "mesh.hpp"
#include "topology.hpp"
#include "util.hpp"
#include "macros.hpp"
#include <vector>
#include <utility>

namespace hms
{


/* this function is mostly for integration testing */
Mesh createRectangularMesh(
	Index  nCellsX,
	Index  nCellsY,
	scalar dX,
	scalar dY,
	scalar offsetX,
	scalar offsetY
){
	return uniMeshAs<Mesh>(
		nCellsX, nCellsY,
		dX, dY, true,
		offsetX, offsetY
	);
}



/**
 * @brief Each edge belongs to one or two domain cells. Once the cell loops are
 * known, the adjacent cells for each edge can be found
 * @param domainCellEdges The edge loops which constitute domain cells
 * @param nDomainCellEdges The number of edges for each cell
 * @param nEdges The number of edges in the mesh
 * @return a 2 x n matrix holding the parent domain cells of each edge as
 * columns. Ghost cells get the signal code for uninitialised values.
 */
auto findAdjacentCells(
	const Ref<const ArrayX<char>>& nDomainCellEdges,
	const Ref<const Array4Xi>& domainCellEdges,
	Index nEdges
)-> Array2Xi
{
	Array2Xi adjacentCells {
		Array2Xi::Constant(2, nEdges, signalCode::uninitialised)
	};
	
	/* not parallelisable like this. Dependency on previous iterations */
	for ( Index cell=0; cell<domainCellEdges.cols(); ++cell ){
		for ( char edge {0}; edge<nDomainCellEdges[cell]; ++edge ){
			Array2Xi::ColXpr col {
				adjacentCells.col( domainCellEdges(edge, cell) )
			};
			if (col[0] == signalCode::uninitialised)
				col[0] = cell;
			else if (col[1] == signalCode::uninitialised)
				col[1] = cell;
		}
	}
	
	return adjacentCells;
}



/**
 * @brief Helper function to check whether an edge is a boundary edge, according
 * to its adjacent cells
 * @param adjacentCells The adjacent cells for all edges
 * @param edge The edge index (i.e. which column within adjacentCells)
 * @return true, if the edge has only one initialised parent cell, false if it
 * has two.
 */
bool hasEmptyParent( const Ref<const Array2Xi>& adjacentCells, Index edge )
{
	assert( adjacentCells.col(edge)[0] != signalCode::uninitialised &&
	"Edge must have at least one parent cell to check if it's a boundary edge!"
	);
	
	/* if edge has only one parent domain cell, it's a boundary cell */
	return adjacentCells.col(edge)[1] == signalCode::uninitialised;
}


/**
 * @brief Collects boundary edge indices and returns them as a 
 * std::vector<Index>
 * @param adjacentCells 
 * @return 
 */
auto findEdgesWEmptyParent( const Ref<const Array2Xi>& adjacentCells ) -> IdxVector
{
	IdxVector boundaryEdges;
	
	for ( Index edge {0}; edge < adjacentCells.cols(); ++edge ){
		if ( hasEmptyParent( adjacentCells, edge ) )
			boundaryEdges.push_back(edge);
	}
	return boundaryEdges;
}

/**
 * @brief Sets ghost indices as adjacent cells of boundary edges.
 * 
 * @param adjacentCells The array holding the adjacent cell indices (Edges::m_cells).
 * @param nDomainCells The number of domain cells in the mesh.
 * @param nBoundaryEdges The number of boundary edges in the mesh.
 */
void setGhostsAsAdjacent(
	Array2Xi& adjacentCells, Index nDomainCells, Index nBoundaryEdges
){
	adjacentCells.row(1).tail(nBoundaryEdges).setLinSpaced(
		nBoundaryEdges, nDomainCells, nDomainCells + nBoundaryEdges - 1
	);
}

/**
 * @brief Swaps the vertex and adjacent cell indices of two edges.
 * 
 * @param edgeA 
 * @param edgeB 
 * @param edgeVertexIndices 
 * @param edgeAdjacentCells 
 */
void swapEdges(
	Index edgeA,
	Index edgeB,
	Array2Xi& edgeVertexIndices,
	Array2Xi& edgeAdjacentCells
){
	swapCols(edgeA, edgeB, edgeVertexIndices, edgeAdjacentCells);
}

/**
 * @brief Swaps the vertex and adjacent cell indices of two edges and also
 * swaps their indices in an array holding the domain cell edges.
 * 
 * @param edgeA 
 * @param edgeB 
 * @param edgeVertexIndices 
 * @param edgeAdjacentCells 
 * @param domainCellEdges 
 */
void swapEdges(
	Index edgeA,
	Index edgeB,
	Array2Xi& edgeVertexIndices,
	Array2Xi& edgeAdjacentCells,
	Array4Xi& domainCellEdges
){
	swapEdges(edgeA, edgeB, edgeVertexIndices, edgeAdjacentCells);

	using colIter = Array4Xi::ColXpr::iterator;
	const Index& nCells { domainCellEdges.cols() };

	/* find edge index in cell loop. no checks, so test carefully! */
	auto findIndexIter = []( Array4Xi::ColXpr cellEdges, Index edge )
		-> colIter
	{
		return std::find( cellEdges.begin(), cellEdges.end(), edge );
	};

	/* find edgeB where edgeA should go and vice versa */
	auto findIters = [&](Index a, Index b)
		-> std::vector<colIter>
	{
		std::vector<colIter> iters;
		for ( int i{0}; i<2; ++i ){
			Index parentCell { edgeAdjacentCells.col(a)[i] };
			if ( parentCell != signalCode::uninitialised && parentCell < nCells ){
				iters.emplace_back( findIndexIter( 
					domainCellEdges.col(parentCell), b
				) );
			}
		}
		return iters;
	};

	/* the edges to be swapped have in total up to four domain cell parents */
	std::vector<colIter>
		edgeAIters { findIters(edgeA,edgeB) },
		edgeBIters { findIters(edgeB,edgeA) };
	
	auto writeNewEdgeIndex = [](Index edge, const std::vector<colIter>& iters){
		for ( auto iter : iters )
			*iter = edge;
	};

	writeNewEdgeIndex(edgeA, edgeAIters);
	writeNewEdgeIndex(edgeB, edgeBIters);
}


auto edgesToSwap(
	const IdxVector& boundaryEdges,
	const Ref<const Array2Xi>& parentCells,
	Index nEdges,
	Index nInnerEdges
) -> Array2Xi {
	
	assert( std::is_sorted( boundaryEdges.begin(), boundaryEdges.end() ) );
	/* boundaryEdges is sorted, so the number of boundary edges that must be
	 * swapped can be determined by finding the first edge index higher than the
	 * number of inner edges. All edges before that one must be swapped, 
	 * whereas all edge indices after that one are larger and therefore already
	 * occupy a boundary edge position. */
	Index nEdgesToSwap { std::lower_bound(
		boundaryEdges.begin(), boundaryEdges.end(), nInnerEdges ) -
		boundaryEdges.begin()
	};
	
	Array2Xi toSwap (2, nEdgesToSwap );
	
	/* start at the first edge index that should represent a boundary edge and
	 * check whether that or a higher one represents an inner edge */
	Index j {0};
	for ( Index i {nInnerEdges}; i<nEdges; ++i ){
		if ( !hasEmptyParent(parentCells, i) ){
			// innerEdgesToSwap.push_back(i);
			toSwap.col(j) = Array2i{ boundaryEdges[j], i };
			++j;
		}
	}

	return toSwap;
}


Mesh edgesAndCellsToMesh(
	Matrix2Xs&& vertexCoords,
	Array2Xi&& edgeVertexIndices,
	ArrayX<char>&& nDomainCellEdges,
	Array4Xi&&      domainCellEdges,
	const Array4Xi& domainCellVertices
){
	Index nDomainCells { nDomainCellEdges.size() };
	Index nEdges { edgeVertexIndices.cols() };

	Array2Xi adjacentCells { findAdjacentCells(
		nDomainCellEdges, domainCellEdges, nEdges
	) };

	/* boundary edges must go to the back */
	/* to swap inner with boundary edges, we first need to find the boundary
	 * edges (by checking their parent cells) */
	IdxVector boundaryEdges { findEdgesWEmptyParent(adjacentCells) };
	assert( std::is_sorted( boundaryEdges.begin(), boundaryEdges.end() ) );

	Index nBoundaryEdges { static_cast<Index>( boundaryEdges.size() ) };
	
	/* whether or not an edge is a boundary edge is to be decided only by its index,
	 * therefore the number of inner edges must later be stored in the mesh */
	Index nInnerEdges { nEdges - nBoundaryEdges };

	/* now we figure out which edges must be swapped (spoiler: it's the actual
	 * boundary edges which appear at an index below nInnerEdges) ... */
	Array2Xi toSwap { edgesToSwap(
		boundaryEdges, adjacentCells, nEdges, nInnerEdges
	) };

	/* ... and swap them and the associated data */
	for ( Array2Xi::ColXpr col : toSwap.colwise() )
		swapEdges(
			col[0], col[1],
			edgeVertexIndices, adjacentCells,
			domainCellEdges
		);
	
	/* now we finish the adjacency table by setting ghost cell indices as
	 * parent cells for boundary edges, i.e. replace the uninitialised values
	 */
	setGhostsAsAdjacent( adjacentCells, nDomainCells, nBoundaryEdges );

	/* with all that information, the mesh can be built */
	return {
		std::move(vertexCoords),
		nInnerEdges,
		std::move(edgeVertexIndices),
		std::move(adjacentCells),
		std::move(nDomainCellEdges),
		std::move( domainCellEdges),
		domainCellVertices
	};
}


Mesh cellVerticesToMesh( /* rename to cellVerticesToMesh */
	Matrix2Xs&& vertexCoords,
	ArrayX<char>&& nDomainCellVertices,
	const Array4Xi& domainCellVertices
){
	Index nDomainCells { nDomainCellVertices.size() };

	Array4Xi domainCellEdges (4, nDomainCells);

	std::vector<IdxVector> edgesAtVertices ( vertexCoords.cols() );
	Array2Xi vertexIndices ( 2, 4 * nDomainCells );
	
	Index edgeIndex {0};
	for ( Index j{0}; j<nDomainCells; ++j ){
		for ( Index i{0}; i<nDomainCellVertices[j]; ++i ){
			/* assume a new edge (block expression doesn't work because of
			 * wrapping around) */
			Array2i edge {
				domainCellVertices.col(j)[ i],
				domainCellVertices.col(j)[(i+1) % nDomainCellVertices[j] ]
			};
			/* now check whether that edge already exists, otherwise store it */
			bool edgeIsNew {true};
			for ( Index e : edgesAtVertices[ edge[0] ] ){
				if ( hms::sameEntries(vertexIndices.col(e), edge) ){
					domainCellEdges(i, j) = e;
					edgeIsNew = false;
					break;
				}
			}
			if (edgeIsNew){
				vertexIndices.col(edgeIndex) = edge;
				edgesAtVertices[ edge[0] ].push_back(edgeIndex);
				edgesAtVertices[ edge[1] ].push_back(edgeIndex);
				domainCellEdges(i, j) = edgeIndex;
				++edgeIndex;
			}
		}
	}
	vertexIndices.conservativeResize(Eigen::NoChange, edgeIndex);
	
	return edgesAndCellsToMesh(
		std::move(vertexCoords),
		std::move(vertexIndices),
		std::move(nDomainCellVertices),
		std::move(domainCellEdges),
		domainCellVertices
	);
}



/**
 * @brief Uses findLoops() on each edge to determine domain cell edges.
 * @return a std::pair containing 1. a matrix whose columns are the edge indices
 * for each cell and 2. a vector, whose entries are the number of edges/vertices
 * for each cell. The cells are ordered by lowest constituting edge index.
 */
auto findDomainCellEdges(
	Index nVertices,
	const Array2Xi& edgeVertices
) -> std::pair<ArrayX<char>, Array4Xi> {

	Index nEdges { edgeVertices.cols() };
	std::vector<IdxVector> edgesAtVertex {
		edgesAtVertices( nVertices, edgeVertices )
	};
	
	/* When the boundary edges n_b and number of internal cells (n_t for 
	 * triangular and n_q for quadrilateral) cells are known, the number of
	 * internal edges n_i is n_i = 1/2 * ( 3*n_t + 4*n_q - n_b )
	 * Therefore, when assuming only triangular cells (smaller cells/edges ratio
	 * -> safe side) and ignoring the number of boundary cells (safe for large
	 * meshes), then the number of domain cells is approximately 2/3 * (n_i+n_b)
	 * */
	std::vector<Array4i> domainCellEdges;
	std::vector<char>   nDomainCellEdges;
	 domainCellEdges.reserve( 2 * nEdges / 3 );
	nDomainCellEdges.reserve( domainCellEdges.capacity() );
	
	#ifdef _OPENMP
	/* parallel version */
	#pragma omp parallel
	{
		decltype(domainCellEdges)   domainCellEdges_perThread;
		decltype(nDomainCellEdges) nDomainCellEdges_perThread;
		domainCellEdges_perThread.reserve(
			domainCellEdges.capacity() / omp_get_num_threads() );
		nDomainCellEdges_perThread.reserve(
			domainCellEdges.capacity() / omp_get_num_threads() );
		
		#pragma omp for schedule(static)
		for ( Index i=0; i < nEdges; ++i )
		{	
			/* find loops for current edge, then store parent cell for each edge
			 * of each loop found. If a found loop constitutes a new cell, store
			 * that cell */
			for ( IdxVector loop : findLoops( edgeVertices, i, edgesAtVertex ) )
			{
				domainCellEdges_perThread.emplace_back(
					hms::toEigen<Array4i>(loop)
				);
				nDomainCellEdges_perThread.push_back(
					static_cast<char>( loop.size() )
				);
			}
		}
		/* make sure cells are written in order */
		for ( Index i=0; i<omp_get_num_threads(); ++i ){
			#pragma omp barrier
			if ( i==omp_get_thread_num() ){
				domainCellEdges.insert( domainCellEdges.end(),
					std::make_move_iterator( domainCellEdges_perThread.begin() ),
					std::make_move_iterator( domainCellEdges_perThread.end() )
				);
				nDomainCellEdges.insert( nDomainCellEdges.end(),
					std::make_move_iterator( nDomainCellEdges_perThread.begin() ),
					std::make_move_iterator( nDomainCellEdges_perThread.end() )
				);
			}
		}
	}
	
	#else
	/* sequential version */
	for ( Index i=0; i < nEdges; ++i )
	{	
		/* find loops for current edge, then store parent cell for each edge
		 * of each loop found. If a found loop constitutes a new cell, store
		 * that cell */
		for ( IdxVector loop : findLoops( edgeVertices, i, edgesAtVertex ) )
		{
			 domainCellEdges.emplace_back( hms::toEigen<Array4i>(loop) );
			nDomainCellEdges.push_back( static_cast<char>( loop.size() ) );
		}
	}
	#endif
	
	return {
		hms::toEigen<ArrayX<char>>(nDomainCellEdges),
		hms::toEigen<Array4Xi>(domainCellEdges)
	};
}

/** 
 * @overload
 */
auto findDomainCellEdges(const Mesh& mesh) 
	-> std::pair<ArrayX<char>, Array4Xi>
{
	return findDomainCellEdges( mesh.nVertices(), mesh.edges().vertices() );
}

Mesh edgesToMesh(
	Matrix2Xs&& vertexCoords,
	Array2Xi&& edgeVertices
){

	Array4Xi domainCellEdges;
	ArrayX<char> nDomainCellEdges;
	std::tie(nDomainCellEdges, domainCellEdges) = findDomainCellEdges(
		vertexCoords.cols(), edgeVertices
	);

	return edgesAndCellsToMesh(
		std::move(vertexCoords),
		std::move(edgeVertices),
		std::move(nDomainCellEdges),
		std::move(domainCellEdges)
	);
}


} // end namespace hms